import math, sys

g_RARE = "_RARE_"

def read_counts_file(filename):
    counts = {}
    bigram_counts = {}
    f = open(filename)
    for line in f:
        tokens = line.split()
        if len(tokens) <> 4:
            continue
        if tokens[1] == "WORDTAG":
            count = int(tokens[0])
            tag = tokens[2]
            word = tokens[3]
            counts[(tag, word)] = count
        elif tokens[1] == "2-GRAM":
            count = int(tokens[0])
            first = tokens[2]
            second = tokens[3]
            bigram_counts[(first, second)] = count
    f.close()
    return counts, bigram_counts

def compute_prob(counts):
    counts = collapse_rare_words(counts)
    prob = {}
    tag_counts = {}
    for pair, count in counts.viewitems():
        tag = pair[0]
        accum(tag_counts, tag, count)
    for pair, count in counts.viewitems():
        tag = pair[0]
        word = pair[1]
        tag_count = tag_counts[tag]
        dict2_add(prob, word, tag, count * 1.0 / tag_count)
    return prob

def compute_bigram_prob(counts):
    first_counts = {}
    prob = {}
    for pair, count in counts.viewitems():
        first = pair[0]
        accum(first_counts, first, count)
    for pair, count in counts.viewitems():
        first, second = pair
        first_count = first_counts[first]
        dict2_add(prob, first, second, count * 1.0 / first_count)
    return prob

def collapse_rare_words(counts):
    word_counts = {}
    for pair, count in counts.viewitems():
        word = pair[1]
        accum(word_counts, word, count)
    rare_words = set([word for word, count in word_counts.viewitems() if count < 5])
    new_counts = {}
    for pair, count in counts.viewitems():
        tag = pair[0]
        word = pair[1]
        if word in rare_words:
            word = g_RARE
        accum(new_counts, (tag, word), count)
    return new_counts

def accum(d, key, value):
    d[key] = d.get(key, 0) + value

def dict2_add(d, key1, key2, value):
    d2 = d.get(key1, {})
    d2[key2] = value
    d[key1] = d2

def bigram_tag_file(sent_file, prob, bigram_prob):
    f = open(sent_file)
    sent = []
    for line in f:
        tokens = line.split()
        if len(tokens) <> 1:
            bigram_tag(sent, prob, bigram_prob)
            sent = []
            continue
        word = tokens[0]
        sent += [word]
    f.close

def get_states(bigram_prob):
    return [s for s in bigram_prob.viewkeys() if s <> "*" and s <> "STOP"]
    
def get_words(sent):
    words = [w for w in set(sent)]
    word_idx = dict([(w, i) for i, w in enumerate(words)])
    return words, word_idx

def replace_by_rare(sent, prob):
    return [w if w in prob else g_RARE for w in sent]

def init_list(m, n):
    return [[0] * n for x in range(m)]

def argmax_max(ls):
    return max(enumerate(ls), key = lambda x : x[1])

def bigram_tag(sent, prob, bigram_prob):
    states = get_states(bigram_prob)
    N = len(states)
    origin_sent = sent
    sent = replace_by_rare(sent, prob)
    words, word_idx = get_words(sent)
    sent = [word_idx[w] for w in sent]
    W = len(words)
    prob_b = init_list(N, W)
    for w in range(W):
        p = prob.get(words[w], {})
        for s in range(N):
            prob_b[s][w] = p.get(states[s], 0)
    def lookup(s1, s2):
        return bigram_prob.get(s1, {}).get(s2, 0)
    prob_a = init_list(N, N)
    for s1 in range(N):
        for s2 in range(N):
            prob_a[s1][s2] = lookup(states[s1], states[s2])
    T = len(sent)
    v = init_list(N, T)
    bv = init_list(N, T)
    for s in range(N):
        v[s][0] = lookup("*", states[s]) * prob_b[s][sent[0]]
        bv[s][0] = 0
    for t in range(1, T):
        for s in range(N):
            bv[s][t], v[s][t] = argmax_max([v[s1][t-1] * prob_a[s1][s] for s1 in range(N)])
            v[s][t] *= prob_b[s][sent[t]]
    bvf, vf = argmax_max([v[s1][T-1] * lookup(states[s1], "STOP") for s1 in range(N)])
    tags = [0] * T
    scores = [0.0] * T
    tags[T - 1] = bvf
    scores[T - 1] = v[tags[T - 1]][T - 1]
    for t in range(T-2, -1, -1):
        tags[t] = bv[tags[t+1]][t+1]
        scores[t] = v[tags[t]][t]
    tags = [states[s] for s in tags];
    for t in range(T):
        print origin_sent[t], tags[t], math.log(scores[t], 2)
    print

counts_file = "ner.counts"
sent_file = "ner_dev.dat"

if len(sys.argv) >= 3:
    counts_file = sys.argv[1]
    sent_file = sys.argv[2]

counts, bigram_counts = read_counts_file(counts_file)
# print bigram_counts
prob = compute_prob(counts)
bigram_prob = compute_bigram_prob(bigram_counts)
# print prob
bigram_tag_file(sent_file, prob, bigram_prob)
